{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from environment import Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 1：Prepare data (batch)\n",
    "\n",
    "Step 2：Neural Network \n",
    "\n",
    "Step 3: Define loss function and optimizer\n",
    "\n",
    "Step 4：Train the network\n",
    "\n",
    "Step 5：GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1. Prepare Data\n",
    "* Tensor\n",
    "* Shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.02577376365661621,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 32,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 170498071,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "046e3a7a8d794399bde8b8cf05c80713",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/170498071 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data/cifar-10-python.tar.gz to ./data\n",
      "original dtype: <class 'numpy.ndarray'>\n",
      "original shape: (50000, 32, 32, 3)\n"
     ]
    }
   ],
   "source": [
    "# Prepare Data\n",
    "trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)\n",
    "\n",
    "print('original dtype:',type(trainset.data))\n",
    "print('original shape:',trainset.data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch(data,batch_size):\n",
    "    minibatch = random.sample(data, batch_size)\n",
    "    minibatch = np.array(minibatch).transpose(0,3,1,2)\n",
    "    minibatch = torch.tensor(minibatch/ 255.0)\n",
    "    return minibatch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 3, 32, 32])\n",
      "torch.float64\n"
     ]
    }
   ],
   "source": [
    "minibatch = batch(list(trainset.data),32)\n",
    "print(minibatch.shape)\n",
    "print(minibatch.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (210, 160, 3)\n",
      "dtype: uint8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A.L.E: Arcade Learning Environment (version 0.7.5+db37282)\n",
      "[Powered by Stella]\n"
     ]
    }
   ],
   "source": [
    "# what does env.reset() look like?\n",
    "env = Environment('BreakoutNoFrameskip-v4', \"\", atari_wrapper=False, test=False)\n",
    "frame = env.reset()\n",
    "print(\"shape:\",frame.shape)\n",
    "print(\"dtype:\",frame.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (84, 84, 4)\n",
      "dtype: uint8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/yli15/ENTER/lib/python3.9/site-packages/gym/core.py:317: DeprecationWarning: \u001b[33mWARN: Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001b[0m\n",
      "  deprecation(\n",
      "/Users/yli15/ENTER/lib/python3.9/site-packages/gym/utils/passive_env_checker.py:227: DeprecationWarning: \u001b[33mWARN: Core environment is written in old step API which returns one bool instead of two. It is recommended to rewrite the environment with new step API. \u001b[0m\n",
      "  logger.deprecation(\n"
     ]
    }
   ],
   "source": [
    "# what does env.reset() look like?\n",
    "env = Environment('BreakoutNoFrameskip-v4', \"\", atari_wrapper=True, test=False)\n",
    "frame = env.reset()\n",
    "print(\"shape:\",frame.shape)\n",
    "print(\"dtype:\",frame.dtype) \n",
    "# convert to float and between 0 and 1 before training\n",
    "# convert to tensor, shape as [batch_size, channel, H, W]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2. Neural Network\n",
    "* init()\n",
    "  * layers\n",
    "* forward\n",
    "  * activation function + combine layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        \"\"\"\n",
    "        Define layers\n",
    "        \"\"\"\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Add activation function\n",
    "        Combine layers to a Neural Network\n",
    "        \"\"\"\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3. Define loss and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4. Train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(),\n",
    "     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,\n",
    "                                          shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,  2000] loss: 2.216\n",
      "[1,  4000] loss: 1.845\n",
      "[1,  6000] loss: 1.666\n",
      "[1,  8000] loss: 1.570\n",
      "[1, 10000] loss: 1.521\n",
      "[1, 12000] loss: 1.465\n",
      "[2,  2000] loss: 1.365\n",
      "[2,  4000] loss: 1.351\n",
      "[2,  6000] loss: 1.335\n",
      "[2,  8000] loss: 1.327\n",
      "[2, 10000] loss: 1.275\n",
      "[2, 12000] loss: 1.292\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(2):  # loop over the dataset multiple times\n",
    "\n",
    "    running_loss = 0.0\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        ####################################################\n",
    "        # tensor\n",
    "        # batch_size, channel, H, W\n",
    "        inputs, labels = data\n",
    "\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        ####################################################\n",
    "        \n",
    "        # print statistics\n",
    "        running_loss += loss.item()\n",
    "        if i % 2000 == 1999:    # print every 2000 mini-batches\n",
    "            print('[%d, %5d] loss: %.3f' %\n",
    "                  (epoch + 1, i + 1, running_loss / 2000))\n",
    "            running_loss = 0.0\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5. GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Assuming that we are on a CUDA machine, this should print a CUDA device:\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign network to cuda\n",
    "net = Net().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,  2000] loss: 2.304\n",
      "[1,  4000] loss: 2.305\n",
      "[1,  6000] loss: 2.305\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# Assign input and output to cuda\n",
    "for epoch in range(2):  # loop over the dataset multiple times\n",
    "\n",
    "    running_loss = 0.0\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        ####################################################\n",
    "        # tensor\n",
    "        # batch_size, channel, H, W\n",
    "        inputs, labels = data[0].to(device), data[1].to(device)\n",
    "\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        ####################################################\n",
    "        \n",
    "        # print statistics\n",
    "        running_loss += loss.item()\n",
    "        if i % 2000 == 1999:    # print every 2000 mini-batches\n",
    "            print('[%d, %5d] loss: %.3f' %\n",
    "                  (epoch + 1, i + 1, running_loss / 2000))\n",
    "            running_loss = 0.0\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
